#include "pasex.hh"
#include "parse.hh"
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include <queue>
#include <stack>
#include <map>
#include <set>

#define ENABLE_DEBUG 1

#if ENABLE_DEBUG
#	define DEBUG(X...) X
#else
#	define DEBUG(X...)
#endif


namespace parse
{
	symbol::symbol(std::string s) : s(s)
	{
		terminal = s[0] != '$' && s[0] != '!';
		mark = s[0] == '@';
		start = s[0] == '!';
	}

	std::istream & operator>>(std::istream & is, symbol & s)
	{
		std::string str;
		is >> str;
		symbol tmp(str);
		s = std::move(tmp);
		return is;
	}

	bool operator<(symbol const & a, symbol const & b) { return a.s < b.s; }
	bool operator==(symbol const & a, symbol const & b) { return a.s == b.s; }


	context_free_rule:: context_free_rule(symbol const& lhs, std::vector<symbol> const& rhs)
		: lhs(lhs), rhs(rhs) {}

	std::istream & operator>>(std::istream & is, context_free_rule & rp)
	{
		symbol lhs;
		is >> rp.lhs;
		for (symbol tmp; is >> tmp; ) rp.rhs.push_back(tmp);
		return is;
	}

	bool operator<(context_free_rule const & a, context_free_rule const & b)
	{
		return a.lhs < b.lhs || (a.lhs == b.lhs && a.rhs < b.rhs);
	}

	bool operator==(context_free_rule const & a, context_free_rule const & b)
	{
		return a.lhs == b.lhs && a.rhs == b.rhs;
	}


	void context_free_grammar::get_first(symbol const& u, symbol const& start)
	{
		vis.insert(u);
		int tu = index(u);
		for (auto i : first_graph[tu]) {
			if (i.terminal) {
				first[start].insert(i);
				continue;
			}
			if (vis.find(i) == vis.end()) get_first(i, start);
		}
	}

	int context_free_grammar::index(symbol const& s)
	{
		if (symbol_index.find(s) != symbol_index.end())
			return symbol_index[s];
		symbol_index[s] = alloc;
		first_graph.push_back(std::vector<symbol>());
		return alloc++;
	}

	void context_free_grammar::add_first_edge(symbol const& u, symbol const& v)
	{
		first_graph[index(u)].push_back(v);
	}

	void context_free_grammar::build_first_graph()
	{
		for (auto r : rps) {
			for (auto i : r.second.rhs) {
				add_first_edge(r.first, i);
				if (i.terminal || !lead_empty(i)) break;
			}
		}
	}

	void context_free_grammar::init_lead_empty()
	{
		auto trps = rps;
		for (auto & i : trps)
			for (auto it = i.second.rhs.begin(); it != i.second.rhs.end(); )
				if (it->s == "&") it = i.second.rhs.erase(it);
				else ++it;
		while (true) {
			symbol tmp;
			bool opt = true;
			for (auto i : trps)
				if (!i.second.rhs.size()) {
					tmp = i.first;
					opt = false;
					break;
				}
			if (opt) break;
			lead_empty_symbol.insert(tmp);
			for (auto it = trps.begin(); it != trps.end(); )
				if (it->first == tmp) it = trps.erase(it);
				else ++it;
			for (auto & i : trps)
				for (auto it = i.second.rhs.begin(); it != i.second.rhs.end(); ) {
					if (*it == tmp) it = i.second.rhs.erase(it);
					else ++it;
				}
		}
	}

	void context_free_grammar::init()
	{
		init_lead_empty();
		build_first_graph();
		for (auto s : all_symbol) {
			vis.clear();
			get_first(s, s);
		}

		DEBUG(
			std::cout << "\n\e[38;5;215mlead empty: \e[0;35m";
			for (auto i : lead_empty_symbol) std::cout << i.s << ' ';
			std::cout << "\e[0m\n";

			for (auto i : all_symbol) {
				std::cout << "first[\e[38;5;215m" << i.s << "\e[0m] = \e[0;35m";
				for (auto j : first[i]) std::cout << j.s << ' ';
				std::cout << "\e[0m\n";
			}
		)

	}

	void context_free_grammar::push(context_free_rule const & rp)
	{
		rps.insert({rp.lhs, rp});
		all_symbol.insert(rp.lhs);
		if (rp.lhs.start && !start_empty) {
			start = rp.lhs;
			start_empty = true;
		}
		if (rp.rhs[0].terminal)
			for (auto ch : rp.rhs[0].s) alpha.insert(ch);
	}

	bool context_free_grammar::lead_empty(symbol const & s) const
	{
		return lead_empty_symbol.find(s) != lead_empty_symbol.end();
	}

	std::set<symbol> context_free_grammar::first_of_symbols(std::vector<symbol> const & sys, std::set<symbol> const & lookahead) const
	{
		std::set<symbol> la;
		for (auto s : sys) {
			if (s.terminal) { la.insert(s); return la; }
			if (first.find(s) != first.end())
				for (auto sf : first.at(s)) la.insert(sf);
			if (!lead_empty(s)) return la;
		}
		for (auto s : lookahead) la.insert(s);
		return la;
	}


	std::istream & operator>>(std::istream & is, context_free_grammar & cfg)
	{
		std::string line;
		while (std::getline(is, line)) {
			if (line.size() == 0) continue;
			std::stringstream buf(line);
			context_free_rule rp;
			buf >> rp;
			cfg.push(rp);
		}
		cfg.start_ex = cfg.start;
		cfg.start_ex.s += '!';
		cfg.push(context_free_rule(cfg.start_ex, {cfg.start}));
		cfg.init();
		return is;
	}

	bool operator<(ex_context_free_rule const & a, ex_context_free_rule const & b)
	{
		return a.r < b.r || (a.r == b.r && a.point < b.point);
	}


	void lr_set_base::add(ex_context_free_rule const & ecfr, std::set<symbol> const & ss)
	{
		for (auto i : ss) lrs[ecfr].insert(i);
	}

	void lr_set_base::add(ex_context_free_rule const & ecfr, symbol const & ch)
	{
		lrs[ecfr].insert(ch);
	}


	lr_set::lr_set(lr_set_base const & lsb, context_free_grammar const & cfg)
	{
		lrs = lsb.lrs;
		std::queue<ex_context_free_rule> q;
		for (auto r : lsb.lrs) q.push(r.first);
		while (!q.empty()) {
			auto now = q.front(); q.pop();
			if (now.point == (int)now.r.rhs.size()) continue;
			if (now.r.rhs[now.point].terminal) continue;
			auto ter = now.r.rhs[now.point];
			auto range = cfg.rps.equal_range(ter);
			auto ts = cfg.first_of_symbols(std::vector<symbol>(now.r.rhs.begin() + now.point + 1, now.r.rhs.end()), lrs[now]);
			for (auto it = range.first; it != range.second; ++it) {
				ex_context_free_rule tmp(it->second, 0);
				if (!in(ts, lrs[tmp])) {
					q.push(tmp);
					for (auto i : ts) lrs[tmp].insert(i);
				}
			}
		}
	}

	bool lr_set::in(std::set<symbol> const& a, std::set<symbol> const& b)
	{
		for (auto i : a) if (b.find(i) == b.end()) return false;
		return true;
	}

	bool operator<(lr_set const & a, lr_set const & b)
	{
		return a.lrs < b.lrs;
	}


	lr_set_dfa::lr_set_dfa(context_free_grammar const & cfg)
	{
		context_free_rule cfr(cfg.start_ex, {cfg.start});
		lr_set_base lsb;
		lsb.add(ex_context_free_rule(cfr, 0), symbol("#"));
		node_index(lr_set(lsb, cfg));
		for (int i = 0; i < (int)all.size(); i++) {
			int id = node_index(all[i]);
			std::map<symbol, lr_set_base> tmp;
			for (auto ls : all[i].lrs) {
				if (ls.first.point == (int)ls.first.r.rhs.size()) continue;
				auto tt = ls.first;
				tt.point++;
				tmp[ls.first.r.rhs[ls.first.point]].add(tt, ls.second);
			}
			for (auto i : tmp)
				add_edge(id, node_index(lr_set(i.second, cfg)), i.first);
		}
	}

	int lr_set_dfa::node_index(lr_set const & ls)
	{
		if (index.find(ls) != index.end()) return index[ls];
		return index[ls] = alloc_node(ls);
	}

	int lr_set_dfa::alloc_node(lr_set const & ls)
	{
		std::map<symbol, int> m;
		graph.push_back(m);
		all.push_back(ls);
		return alloc++;
	}

	bool lr_set_dfa::acc(int state)
	{
		for (auto i : all[state].lrs) {
			auto tmp = i.first.r.lhs.s;
			if (tmp[0] == '!' && tmp[tmp.size() - 1] == '!') return true;
		}
		return false;
	}

	int lr_set_dfa::size() const
	{
		return all.size();
	}

	int lr_set_dfa::sr_conflict_count() const
	{
		int count = 0;
		for (auto i : all) {
			std::set<symbol> now;
			bool conflict = false;
			for (auto j : i.lrs) {
				if (j.first.point != (int)j.first.r.rhs.size()) continue;
				for (auto t : j.second) {
					if (t.s == "&") continue;
					now.insert(t);
				}
			}
			for (auto j : i.lrs) {
				if (j.first.point == (int)j.first.r.rhs.size()) continue;
				for (auto t : j.second) {
					if (now.find(t) != now.end()) {
						conflict = true;
						break;
					}
				}
				if (conflict) break;
			}
			count += conflict;
		}
		return count;
	}

	int lr_set_dfa::rr_conflict_count() const
	{
		int count = 0;
		for (auto i : all) {
			std::set<symbol> now;
			bool conflict = false;
			for (auto j : i.lrs) {
				if (j.first.point != (int)j.first.r.rhs.size()) continue;
				for (auto t : j.second)
					if (now.find(t) != now.end()) {
						conflict = true;
						break;
					}
				if (conflict) break;
				for (auto t : j.second)
					now.insert(t);
			}
			count += conflict;
		}
		return count;
	}

	void print()
	{
	}

	bool lr_set_dfa::parse(std::string const & s, lex::dfa::dfa & d)
	{
		d.lex_init();
		std::string error{"\e[1;31m[FAIL]\e[0;31m some errors occured.\e[0m\n"};
		auto code = s + char(0);
		auto it = code.begin();
		std::stack<int> state;
		std::stack<symbol> symbols;
		state.push(0);
		symbols.push(symbol("#"));
		std::pair<std::string, std::string> token;
		bool reduce = false;
		while (true) {
			if (!reduce) token = d.lex(code, it);
			if (token.first == "#" && token.second == "#") {
				if (state.size() == 2 && acc(state.top())) return true;
			}
			if (token.second == "^") continue;
			auto now = state.top();
			reduce = false;
			symbol tmp(token.first);
			if (graph[now].find(token.first) != graph[now].end()) {
				state.push(graph[now][symbol(token.first)]);
				symbols.push(tmp);
			} else if (graph[now].find(token.second) != graph[now].end()) { // FIXME extend marks
				state.push(graph[now][symbol(token.second)]);
				symbols.push(symbol(token.second));
			} else {
				for (auto i : all[now].lrs) {
					if (i.first.point != (int)i.first.r.rhs.size()) continue;
					bool find = false;
					if (i.second.find(tmp) != i.second.end() || i.second.find(symbol(token.second)) != i.second.end()) { // FIXME extend marksc
						for (int c = i.first.r.rhs.size(); c; c--) {
							if (state.empty() || symbols.empty()) {
								std::cout << error;
								return false;
							}
							state.pop();
							symbols.pop();
						}
						find = true;
					}
					if (!find) continue;
					symbols.push(i.first.r.lhs);
					if (state.empty()) {
						std::cout << error;
						return false;
					}
					if (graph[state.top()].find(symbols.top()) == graph[state.top()].end()) {
						std::cout << error;
						return false;
					}
					state.push(graph[state.top()][symbols.top()]);
					reduce = true;
					break;
				}
				if (!reduce) {
					std::cout << error;
					return false;
				}
			}
		}
		return true;
	}
}

#undef DEBUG
#undef ENABLE_DEBUG

